import torch
import torch.nn as nn
import torchvision
from torch.utils.tensorboard import SummaryWriter
from torchvision import datasets, transforms

from Forecase import Forecase
from sort.data.DataHandler import DataHandler

if __name__ == '__main__':

    # 创建数据集
    data_dir = 'I:\\Workspace\\Dataset\\data_inception\\train'  # 你的数据集路径
    batch_size = 32

    # 数据预处理和增强
    transform = DataHandler(data_dir, batch_size).transform

    train_dataset = datasets.ImageFolder(data_dir, transform=transform)

    # 数据集拆分比例
    train_size = int(0.7 * len(train_dataset))
    val_size = len(train_dataset) - train_size

    # 随机拆分数据集
    train_dataset, val_dataset = torch.utils.data.random_split(train_dataset, [train_size, val_size])

    # 数据加载器
    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=batch_size)

    for X, y in train_loader:
        print(f"Shape of X [N, C, H, W]: {X.shape}")
        print(f"Shape of y: {y.shape} {y.dtype}")
        break

    # 获取 cpu, gpu 或 mps 设备用于加速训练.
    device = (
        "cuda"
        if torch.cuda.is_available()  # 如果GPU可用，则使用GPU
        else "mps"
        if torch.backends.mps.is_available()  # 如果MPS可用，则使用MPS
        else "cpu"  # 否则使用CPU
    )
    # torch.backends.mps.is_available()

    print(f"Using {device} device")

    classes = train_dataset.dataset.classes

    num_classes = len(classes)

    # 加载模型结构和权重
    # model = torch.load('your_model.pth')

    # 或者加载模型权重到预先创建的模型
    # model_path = 'model/ImageSort.pth'
    # model = NeuralNetwork(len(classes))
    # model.load_state_dict(torch.load(model_path))

    #
    model = torchvision.models.resnet50(weights=None)

    num_features = model.fc.in_features
    model.fc = nn.Linear(num_features, num_classes)

    model.load_state_dict(torch.load('model/ResNet50-ImageSort.pth'))

    forecase = Forecase(model, transform)
    forecase.run_forecast('I:\\Workspace\\Dataset\\data_inception\\test\\5', classes)
